from langchain_community.vectorstores import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough


from model import get_model
from Embed_model import BGE_Embed

model = get_model()
embedding = BGE_Embed()

# 加载知识库
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# sore=vectorstore.similarity_search_with_score('成都4级大概多少一个？你给我一个，我算算成本,')
# print(sore)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

# 检索出的文档
result = retriever.invoke('扫码多少钱？')
print(result)
# # print(len(result))
# exit()


def doc_formatter(docs):
    # 使用集合去重，然后再合并成字符串
    unique_docs = set([doc.page_content for doc in docs])
    return '\n\n'.join(unique_docs)

# 此处 {context} 代表的是检索结果
# 提示词和问题加上检索的结果
system_prompt = (
    "你是一个云剪、实探业务的聊天客服"
    "根据检索到的上下文来回答问题。"
    "回答要精简"
    "\n\n"
    "{context}"
)

# 创建提示词模板
prompt = ChatPromptTemplate.from_messages([
    ('system', system_prompt),
    ('user', '{input}')
])


# RunnablePassthrough() 用于允许我们将用户的问题等会再传给prompt和model,在chain中串联时可以动态的设置用户问题
# 代表 agent 的输入，原封不动的会被替换到 RunnablePassthrough()
# StrOutputParser() 字符串输出转换器，将模型输出的 AIMessage 对象转换成字符串
chain = (
        # context是上下文，传的检索器结果，input是用户问题
        {'context': retriever | doc_formatter, 'input': RunnablePassthrough()}
        | prompt
        | model
        | (lambda x: x.content)
        | StrOutputParser()
)

res = chain.invoke('扫码多少钱？')
print(res)


